package dan.frontend

import chisel3._
import chisel3.util._
import dan.utility.FrontendUtility._
import dan.utility.ImplicitCast.uintToBitPat
import chisel3.experimental.treedump

class BiMMeta extends Bundle with HasBiMParam{
    val state = UInt(2.W)
}

class BiMIO extends Bundle with HasBiMParam{
    val s0Valid = Input(Bool())
    val s0PC = Input(UInt(vaBits.W))
    val s2Taken = Output(Vec(fetchWidth, Bool()))
    val s3Meta = Output(Vec(fetchWidth, new BiMMeta()))
    val s1Update = Input(Valid(new PredictionUpdate()))
}

class BiM extends Module with HasBiMParam{
    val io = IO(new BiMIO())
    val states = SyncReadMem(setNum, Vec(fetchWidth, UInt(2.W)))
    val resetEnable = RegInit(false.B)
    val resetIdx = RegInit(0.U(log2Ceil(setNum).W))
    when(resetEnable){
        resetIdx := resetIdx + 1.U
    }
    when(resetIdx === (setNum - 1).U){
        resetEnable := false.B
    }
    // 读出pc对应状态
    val s0Valid = io.s0Valid
    val s0Idx = fetchIdx(io.s0PC)
    val s2States = RegNext(VecInit(states.read(s0Idx.asUInt, s0Valid).map(_.asTypeOf(UInt(2.W)))))
    for (i <- 0 until fetchWidth) {
        io.s2Taken(i) := s2States(i)(1)
    }
    // s2Meta <-> s2States
    val s2Meta = Wire(Vec(fetchWidth, new BiMMeta()))
    for (i <- 0 until fetchWidth){
        s2Meta(i).state := s2States(i)
    }
    io.s3Meta := RegNext(s2Meta)

    val s1Update = io.s1Update
    // 这个fetchIdx真的没问题吗 ??
    val s1UpdateIdx = fetchIdx(s1Update.bits.pc)
    // 旧值
    val s1UpdateMeta = VecInit(s1Update.bits.meta.map(_.bimMeta))
    val s1UpdateMask = Wire(Vec(fetchWidth, Bool()))
    val s1UpdateState = Wire(Vec(fetchWidth, UInt(2.W)))

    for(i <- 0 until fetchWidth){
        // 默认值
        s1UpdateMask(i) := false.B
        s1UpdateState(i) := s1UpdateMeta(i).asUInt
        when(s1Update.valid && (s1Update.bits.brMask(i) ||
            (s1Update.bits.idx.valid && s1Update.bits.idx.bits === i.U)
        )){
            val taken = s1Update.bits.idx.valid && s1Update.bits.idx.bits === i.U && (
                (s1Update.bits.isBr && s1Update.bits.brMask(i) && s1Update.bits.taken) || s1Update.bits.isJal
            )
            s1UpdateMask(i) := true.B
            s1UpdateState(i) := changeState(s1UpdateMeta(i).asUInt, taken)
        }
    }
    states.write(Mux(resetEnable, resetIdx, s1UpdateIdx), 
        Mux(resetEnable, VecInit(Seq.fill(fetchWidth){2.U(2.W)}), s1UpdateState),
        Mux(resetEnable, ~(0.U(fetchWidth.W)), s1UpdateMask.asUInt).asBools
    )
}


